package current.lockfree;

import org.junit.jupiter.api.Test;
import java.util.concurrent.atomic.AtomicStampedReference;
import static org.junit.jupiter.api.Assertions.assertEquals;

public class LockFreeStack<T> {

    AtomicStampedReference<Node<T>> head;

    static class Node<T> {
        Node<T> next;
        T value;

        public Node(T value) {
            this.value = value;
            this.next = null;
        }
    }


    public LockFreeStack() {
        var headNode = new Node<T>(null);
        head = new AtomicStampedReference<>(headNode, 0);
    }

    public void push(T v) {
        // 假如新Node为99
        var newNode = new Node<T>(v);
        while (true) {
            // 获取版本号
            int stamp = head.getStamp();

            // 头插法，当前为98
            var ref = head.getReference();
            // 99的next指向98
            newNode.next = ref;
            // 解决ABA问题
            if (head.compareAndSet(ref, newNode, stamp, stamp + 1)) {
                // 插入成功并正确设置了next就返回退出循环
                return;
            }
        }
    }

    public T pop() {
        while (true) {
            int stamp = head.getStamp();
            // 假如拿到99
            var ref = head.getReference();
            if (ref.next == null) {
                // 直到没有下一节点就退出循环
                return null;
            }
            // 这里拿到98
            var next = ref.next;
            // 头部删除，同样要注意ABA问题
            head.compareAndSet(ref, next, stamp, stamp + 1);
            return ref.value;
        }
    }

    @Test
    public void testSingleThread() {
        LockFreeStack<Integer> stack = new LockFreeStack<>();

        for (int i = 0; i < 100; i++) {
            // 0->99进
            stack.push(i);
        }

        Integer j = null;
        Integer k = 99;
        while ((j = stack.pop()) != null) {
            // 99->0出
            assertEquals(j + "", k-- + "");
        }
    }

    @Test
    public void testMultiThread() throws InterruptedException {
        LockFreeStack<Integer> stack = new LockFreeStack<>();

        for (int i = 0; i < 16; i++) {
            Thread t = new Thread(() -> {
                for (int j = 0; j < 100; j++) {
                    stack.push(j);
                }
            });

            t.start();
            // 确保所有线程走完
            t.join();
        }

        Integer k = 0;
        while (stack.pop() != null) {
            k++;
        }

        assertEquals(k + "", "1600");
    }
}
